#include "codegenerator.h"
#include <sstream>
#include <algorithm>
#include <stdexcept>
#include <iostream>

using namespace std;

string CodeGenerator::CreateLabel() {//生成唯一标签
    return "L" + to_string(++labelCounter);
}

//寄存器分配
string CodeGenerator::Record::allocReg() {
    if (freeRegs.empty()) {
        throw runtime_error("寄存器用完");
    }
    string reg = freeRegs.back();
    freeRegs.pop_back();
    usedRegs.push_back(reg);
    
    //记录调用者保存的寄存器
    if (reg[0] == 't' || reg[0] == 'a') {
        callerSaved.push_back(reg);
    }
    
    return reg;
}
//寄存器释放
void CodeGenerator::Record::freeReg(const string& reg) {
    auto it = find(usedRegs.begin(), usedRegs.end(), reg);
    if (it != usedRegs.end()) {
        usedRegs.erase(it);
        freeRegs.push_back(reg);
        
        //从callerSaved中移除
        auto it2 = find(callerSaved.begin(), callerSaved.end(), reg);
        if (it2 != callerSaved.end()) {
            callerSaved.erase(it2);
        }
    }
}

//表达式生成
string CodeGenerator::genExpr(ostream& os, Record& r, Expr* expr) {
    if (auto e = dynamic_cast<EInt*>(expr)) { //取常量
        string reg = r.allocReg();
        os << "  li " << reg << ", " << e->value << "\n";
        return reg;
    }
    if (auto e = dynamic_cast<EId*>(expr)) { //取变量
        string reg = r.allocReg();
        int offset = r.varMap.at(e->id);
        os << "  lw " << reg << ", " << offset << "(s0)\n";
        return reg;
    }
    if (auto e = dynamic_cast<EUnOp*>(expr)) { //一元运算
        string srcReg = genExpr(os, r, e->expr.get());
        string destReg = r.allocReg();
        if (e->op == UnOp::Neg) {
            os << "  neg " << destReg << ", " << srcReg << "\n";
        } else { // Not
            os << "  seqz " << destReg << ", " << srcReg << "\n";
        }
        r.freeReg(srcReg);
        return destReg;
    }
    if (auto e = dynamic_cast<EBinOp*>(expr)) { //二元运算
        string leftReg = genExpr(os, r, e->left.get());
        string rightReg = genExpr(os, r, e->right.get());
        string destReg = r.allocReg();
        
        switch (e->op) {
            case BinOp::Add: os << "  add " << destReg << ", " << leftReg << ", " << rightReg << "\n"; break;
            case BinOp::Sub: os << "  sub " << destReg << ", " << leftReg << ", " << rightReg << "\n"; break;
            case BinOp::Mul: os << "  mul " << destReg << ", " << leftReg << ", " << rightReg << "\n"; break;
            case BinOp::Div: os << "  div " << destReg << ", " << leftReg << ", " << rightReg << "\n"; break;
            case BinOp::Mod: os << "  rem " << destReg << ", " << leftReg << ", " << rightReg << "\n"; break;
            case BinOp::Eq:  os << "  xor " << destReg << ", " << leftReg << ", " << rightReg << "\n  seqz " << destReg << ", " << destReg << "\n"; break;
            case BinOp::Neq: os << "  xor " << destReg << ", " << leftReg << ", " << rightReg << "\n  snez " << destReg << ", " << destReg << "\n"; break;
            case BinOp::Lt:  os << "  slt " << destReg << ", " << leftReg << ", " << rightReg << "\n"; break;
            case BinOp::Le:  os << "  sgt " << destReg << ", " << leftReg << ", " << rightReg << "\n  xori " << destReg << ", " << destReg << ", 1\n"; break;
            case BinOp::Gt:  os << "  sgt " << destReg << ", " << leftReg << ", " << rightReg << "\n"; break;
            case BinOp::Ge:  os << "  slt " << destReg << ", " << leftReg << ", " << rightReg << "\n  xori " << destReg << ", " << destReg << ", 1\n"; break;
            case BinOp::And: 
                os << "  snez " << leftReg << ", " << leftReg << "\n";
                os << "  snez " << rightReg << ", " << rightReg << "\n";
                os << "  and " << destReg << ", " << leftReg << ", " << rightReg << "\n";
                break;
            case BinOp::Or: 
                os << "  or " << destReg << ", " << leftReg << ", " << rightReg << "\n";
                os << "  snez " << destReg << ", " << destReg << "\n";
                break;
            default: throw runtime_error("不支持的二元操作");
        }
        r.freeReg(leftReg);
        r.freeReg(rightReg);
        return destReg;
    }
    if (auto e = dynamic_cast<ECall*>(expr)) { //函数调用
        //保存调用者保存的寄存器
        vector<string> toSave;
        for (const auto& reg : r.callerSaved) {
            if (find(r.usedRegs.begin(), r.usedRegs.end(), reg) != r.usedRegs.end()) {
                toSave.push_back(reg);
            }
        }
        int saveSpace = toSave.size() * 4;

        //调整栈指针并保存寄存器
        if (saveSpace > 0) {
            os << "  addi sp, sp, -" << saveSpace << "\n";
            for (int i = 0; i < toSave.size(); i++) {
                os << "  sw " << toSave[i] << ", " << i*4 << "(sp)\n";
            }
        }

        //取出参数
        vector<string> argRegs;
        for (int i = 0; i < e->args.size(); i++) {
            string argReg = genExpr(os, r, e->args[i].get());
            argRegs.push_back(argReg);
        }

        //参数存入参数寄存器
        for (int i = 0; i < min(8, (int)argRegs.size()); i++) {
            string target = "a" + to_string(i);
            if (argRegs[i] != target) {
                os << "  mv " << target << ", " << argRegs[i] << "\n";
            }
            r.freeReg(argRegs[i]);
        }
        //多余参数压栈
        for (int i = 8; i < argRegs.size(); i++) {
            //前面已经调整了sp，所以额外的参数应该放在当前调整后的sp之上
            os << "  sw " << argRegs[i] << ", " << saveSpace + (i-8)*4 << "(sp)\n";
            r.freeReg(argRegs[i]);
        }

        //保存返回地址
        string savedRA = r.allocReg();
        os << "  mv " << savedRA << ", ra\n";
        
        //生成调用指令
        os << "  call " << e->id << "\n";
        
        //恢复返回地址
        os << "  mv ra, " << savedRA << "\n";
        r.freeReg(savedRA);
        
        //恢复调用者保存的寄存器
        if (saveSpace > 0) {
            for (int i = 0; i < toSave.size(); i++) {
                os << "  lw " << toSave[i] << ", " << i*4 << "(sp)\n";
            }
            os << "  addi sp, sp, " << saveSpace << "\n";
        }
        
        //处理返回值
        string resultReg = r.allocReg();
        os << "  mv " << resultReg << ", a0\n";
        return resultReg;
    }
    throw runtime_error("未知表达式类型");
}

void CodeGenerator::genStmt(ostream& os, Record& r, Stmt* stmt) {
    if (auto s = dynamic_cast<SExpr*>(stmt)) { //独立表达式
        string reg = genExpr(os, r, s->expr.get());
        r.freeReg(reg);
        return;
    }
    if (auto s = dynamic_cast<SReturn*>(stmt)) { //返回语句
        if (s->expr) {
            string reg = genExpr(os, r, s->expr.get());
            os << "  mv a0, " << reg << "\n";//a0存返回值
            r.freeReg(reg);
        }
        os << "  j " << r.exitLabel << "\n"; //跳转到当前函数退出点标签
        return;
    }
    if (auto s = dynamic_cast<SDecl*>(stmt)) { //声明语句
        string reg = genExpr(os, r, s->expr.get());
        int offset = r.varMap.at(s->id);//SBlock提前分配了栈偏移
        os << "  sw " << reg << ", " << offset << "(s0)\n";
        r.freeReg(reg);
        return;
    }
    if (auto s = dynamic_cast<SAssign*>(stmt)) { //赋值语句
        string reg = genExpr(os, r, s->expr.get());
        int offset = r.varMap.at(s->id);
        os << "  sw " << reg << ", " << offset << "(s0)\n";
        r.freeReg(reg);
        return;
    }
    if (auto s = dynamic_cast<SIf*>(stmt)) { //条件语句
        string elseLabel = CreateLabel();
        string endLabel = CreateLabel();
        
        string condReg = genExpr(os, r, s->cond.get());
        os << "  beqz " << condReg << ", " << elseLabel << "\n";
        r.freeReg(condReg);
        
        genStmt(os, r, s->thenStmt.get());
        
        if (s->elseStmt) {
            os << "  j " << endLabel << "\n";
            os << elseLabel << ":\n";
            genStmt(os, r, s->elseStmt.get());
            os << endLabel << ":\n";
        } else {
            os << elseLabel << ":\n";
        }
        return;
    }
    if (auto s = dynamic_cast<SWhile*>(stmt)) { //循环语句
        LoopContext ctx;
        ctx.startLabel = CreateLabel();
        ctx.endLabel = CreateLabel();
        
        //将新循环压入栈
        r.pushLoop(ctx);
        
        //循环开始标签
        os << ctx.startLabel << ":\n";
        //生成条件表达式
        string condReg = genExpr(os, r, s->cond.get());
        os << "  beqz " << condReg << ", " << ctx.endLabel << "\n";
        r.freeReg(condReg);
        //生成循环体
        genStmt(os, r, s->body.get());
        //跳回循环开始
        os << "  j " << ctx.startLabel << "\n";
        //循环结束标签
        os << ctx.endLabel << ":\n";
        
        //退出当前循环
        r.popLoop();
        return;
    }
    if (auto s = dynamic_cast<SBreak*>(stmt)) { //break语句
        //取栈顶的循环标签，检查是否在循环中
        string endLabel = r.currentEndLabel();
        if (endLabel.empty()) {
            throw runtime_error("break语句不在循环中");
        }
        os << "  j " << endLabel << "\n";
        return;
    }
    if (auto s = dynamic_cast<SContinue*>(stmt)) { //continue语句
        //取栈顶的循环标签，检查是否在循环中
        string startLabel = r.currentStartLabel();
        if (startLabel.empty()) {
            throw runtime_error("continue语句不在循环中");
        }
        os << "  j " << startLabel << "\n";
        return;
    }
    if (auto s = dynamic_cast<SBlock*>(stmt)) { //块语句
        //创建新的作用域（继承父作用域的变量映射）
        Record blockRecord = r;
        
        //为块内的局部变量分配偏移量
        for (auto& st : s->stmts) {
            if (auto decl = dynamic_cast<SDecl*>(st.get())) {
                int offset = blockRecord.currentOffset;
                blockRecord.varMap[decl->id] = offset;
                blockRecord.currentOffset -= 4; //为下一个变量预留空间
            }
        }
        
        //生成块内的语句
        for (auto& st : s->stmts) {
            genStmt(os, blockRecord, st.get());
        }
        return;
    }
    throw runtime_error("未知语句类型");
}

void CodeGenerator::genFunc(ostream& os, FuncDef* func) {
    Record r;
    r.exitLabel = CreateLabel(); //创建函数退出点
    
    //计算栈帧大小
    int numLocals = func->body->countLocalVars();
    int numParams = func->params.size();
    int totalSpace = 8 + 4 * numParams + 4 * numLocals;  // 8(ra+s0) + 参数 + 局部变量
    
    //确保至少有最小栈空间用于函数调用
    if (totalSpace < 128) totalSpace = 128;
    
    //16字节对齐
    totalSpace = (totalSpace + 15) & ~15;
    r.frameSize = totalSpace; //设置帧大小
    
    //函数序言
    os << ".global " << func->name << "\n";
    os << ".type " << func->name << ", @function\n";
    os << func->name << ":\n";
    os << "  addi sp, sp, -" << totalSpace << "\n";
    os << "  sw ra, " << totalSpace - 4 << "(sp)\n";
    os << "  sw s0, " << totalSpace - 8 << "(sp)\n";
    os << "  addi s0, sp, " << totalSpace << "\n";  //s0指向原始sp位置
    
    //初始化当前偏移（用于函数调用时的临时存储）
    r.currentOffset = 0;

    //保存参数到栈帧
    for (int i = 0; i < numParams; i++) {
        int offset = totalSpace - 12 - 4 * i;//offset是存储位置到sp的偏移
        os << "  sw a" << i << ", " << offset << "(sp)\n";
        r.varMap[func->params[i]->id] = offset - totalSpace; //相对于s0的偏移
    }
    
    //设置局部变量偏移：从s0-16开始（跳过ra, s0和参数区域）
    r.currentOffset = -16 - 4 * numParams;
    
    //生成函数体
    genStmt(os, r, func->body.get());
    
    //函数退出点标签
    os << r.exitLabel << ":\n";
    os << "  lw ra, " << totalSpace - 4 << "(sp)\n";
    os << "  lw s0, " << totalSpace - 8 << "(sp)\n";
    os << "  addi sp, sp, " << totalSpace << "\n";
    os << "  ret\n";
    os << ".size " << func->name << ", .-" << func->name << "\n\n";
}

void CodeGenerator::genCompUnit(ostream& os, CompUnit* cu) {
    os << ".text\n";
    os << ".globl main\n"; //确保main函数是全局的
    
    //遍历生成所有函数
    for (auto& func : cu->funcs) {
        genFunc(os, func.get());
    }
}